package com.lxg.ai.config;

import com.lxg.ai.service.RagChatAssistant;
import com.lxg.ai.store.PersistentChatMemoryStore;
import com.lxg.ai.tool.MyWebSearchFactoryTool;
import com.lxg.ai.transformer.MyQueryTransformer;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.tool.DefaultToolExecutor;
import dev.langchain4j.service.tool.ToolProviderResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.google.customsearch.GoogleCustomWebSearchEngine;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.lang.reflect.Method;


@Configuration
public class RagConfig {


    @Autowired
    private ThreadPoolTaskExecutor myAsync;


    @Bean
    public QdrantEmbeddingStore embeddingStore(){
        QdrantEmbeddingStore embeddingStore = QdrantEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6334)
                .collectionName("测试向量集合")
                .build();
        return embeddingStore;
    }

    @Bean
    public RagChatAssistant ragChatAssistant(ChatLanguageModel chatLanguageModel, EmbeddingStore<TextSegment> embeddingStore, EmbeddingModel embeddingModel, ThreadPoolTaskExecutor myAsync) {

        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(3)   // 返回最多3条结果
                .dynamicMaxResults(query -> 3)  // 根据查询动态指定maxResults
//                .minScore(0.75)  // 过滤分数低于0.75的内容
                .dynamicMinScore(query -> 0.75)  // 根据查询动态指定minScore
//                .filter(metadataKey("uuid").isEqualTo("b79c995a-fefd-45b5-bc86-1c92181c75b2"))  // 过滤指定用户的内容
//                .dynamicFilter(query -> {
//                    String userId = String.valueOf(query.metadata().chatMemoryId()) ;
//                    return metadataKey("userId").isEqualTo(userId);  // 动态过滤
//                })
                .build();


        DefaultRetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
//                .queryTransformer(new CompressingQueryTransformer(chatLanguageModel, PromptTemplate.from("")))  // 查询增强 提示词等 具体看源码
                .queryTransformer(new MyQueryTransformer(chatLanguageModel))  // 查询增强 提示词等 具体看源码
                .contentRetriever(contentRetriever) // 内容源 单个直接配置
//                .queryRouter(new DefaultQueryRouter())// 多个内容源，路由
//                .contentAggregator(new DefaultContentAggregator()) // 匹配结果聚合 多个结果选择
//                .contentInjector(new DefaultContentInjector())   // 结果提示词注入
                .executor(myAsync)  // 并行化
                .build();

//        GoogleCustomWebSearchEngine searchEngine = GoogleCustomWebSearchEngine  谷歌的
//                .builder()
//                .apiKey("fdsfsadfsdafsadfsdaf")
//                .csi("fdsfdsafsadfsdafsadf")
//                .logRequests(true)
//                .logResponses(true)
//                .timeout(Duration.ofSeconds(60))
//                .build();

//        SearchApiWebSearchEngine searchEngine = SearchApiWebSearchEngine.builder()
//                .apiKey("fdsafdsaffdsafdsa")// 测试使用
//                .engine("google")
//                .build();

//        WebSearchTool webSearchTool = WebSearchTool.from(searchEngine);

//         业务逻辑 ToolExecutor
        return AiServices.builder(RagChatAssistant.class)
                .chatLanguageModel(chatLanguageModel)
//                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(20))//取最新的10条做上下文//如需自定义持久化方式,可以实现ChatMemoryStore接口:
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder()
                        .id(memoryId)
                        .maxMessages(10)
                        .chatMemoryStore(new PersistentChatMemoryStore())
                        .build())//取最新的10条做上下文//如需自定义持久化方式,可以实现ChatMemoryStore接口:
                .retrievalAugmentor(retrievalAugmentor)
//                .toolProvider(s->{
//                    ToolProviderResult.Builder builder = ToolProviderResult.builder();
//                    if (s.userMessage().singleText().contains("booking")) {
//                        Method method = MyWebSearchFactoryTool.getMethod(webSearchTool);
//                        return builder.add(ToolSpecifications.toolSpecificationFrom(method), new DefaultToolExecutor(webSearchTool, method)).build();
//                    }else {
//                        return null;
//                    }
//
//                })
//                .tools(webSearchTool)
                .build();
    }
}
